import os
import glob
import random
import math
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
import cv2
from tqdm import tqdm
import torch.optim as optim


# ------------------------- 模型定义 -------------------------
def make_net(nplanes_in, kernels, features, bns, acts, dilats, bn_momentum=0.1, padding=1):
    layers = []
    depth = len(features)

    for i in range(depth):
        in_feats = nplanes_in if i == 0 else features[i - 1]

        # 卷积层
        conv = nn.Conv2d(in_feats, features[i],
                         kernel_size=kernels[i],
                         dilation=dilats[i],
                         padding=padding,
                         bias=not bns[i])

        # 初始化
        n = kernels[i] * kernels[i] * features[i]
        conv.weight.data.normal_(0, math.sqrt(2. / n))
        layers.append(conv)

        # 批量归一化
        if bns[i]:
            bn = nn.BatchNorm2d(features[i], momentum=bn_momentum)
            bn.weight.data.normal_(0, math.sqrt(2. / n))
            bn.bias.data.zero_()
            layers.append(bn)

        # 激活函数
        if acts[i] == 'relu':
            layers.append(nn.ReLU(inplace=True))
        elif acts[i] == 'sigmoid':
            layers.append(nn.Sigmoid())

    return nn.Sequential(*layers)


# ------------------------- 数据加载 -------------------------
class TamperDataset(Dataset):
    def __init__(self, tampered_dir, mask_dir, patch_size=256):
        """初始化数据集，自动配对.jpg/.png篡改图像和.png掩模"""
        # 获取所有图像路径
        self.tampered_images = []
        for ext in ['jpg', 'jpeg', 'png']:
            self.tampered_images.extend(glob.glob(os.path.join(tampered_dir, f'*.{ext}')))

        # 验证配对
        self.valid_pairs = []
        for img_path in self.tampered_images:
            base_name = os.path.splitext(os.path.basename(img_path))[0]
            mask_path = os.path.join(mask_dir, f"{base_name}.png")
            if os.path.exists(mask_path):
                self.valid_pairs.append((img_path, mask_path))
            else:
                print(f"Warning: Missing mask for {img_path}")

        self.patch_size = patch_size
        self.transform = transforms.Compose([
            transforms.ToTensor(),
        ])

    def __len__(self):
        return len(self.valid_pairs)

    def __getitem__(self, idx):
        img_path, mask_path = self.valid_pairs[idx]

        # 加载图像和掩模
        img = Image.open(img_path).convert('RGB')
        mask = Image.open(mask_path).convert('L')

        # 数据增强
        if random.random() > 0.5:
            img = transforms.functional.hflip(img)
            mask = transforms.functional.hflip(mask)

        # 随机裁剪
        i, j, h, w = transforms.RandomCrop.get_params(
            img, output_size=(self.patch_size, self.patch_size))

        img_patch = transforms.functional.crop(img, i, j, h, w)
        mask_patch = transforms.functional.crop(mask, i, j, h, w)

        # 转换格式
        img_tensor = self.transform(img_patch)
        mask_tensor = torch.from_numpy(np.array(mask_patch)).float() / 255.0

        return img_tensor, mask_tensor.unsqueeze(0)


# ------------------------- 损失函数 -------------------------
class DiceLoss(nn.Module):
    def forward(self, pred, target):
        smooth = 1.0
        intersection = (pred * target).sum()
        union = pred.sum() + target.sum()
        dice = (2. * intersection + smooth) / (union + smooth)
        return 1 - dice


class TamperLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.bce = nn.BCELoss()
        self.dice = DiceLoss()

    def forward(self, pred, target):
        return self.bce(pred, target) + self.dice(pred, target)


# ------------------------- 训练函数 -------------------------
def train_model(model, train_loader, val_loader, epochs=50, lr=1e-4):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)

    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3)
    criterion = TamperLoss()

    best_val_loss = float('inf')

    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0
        pbar = tqdm(train_loader, desc=f'Epoch {epoch + 1}/{epochs}')
        for img, mask in pbar:
            img, mask = img.to(device), mask.to(device)

            optimizer.zero_grad()
            pred = model(img)
            loss = criterion(pred, mask)
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            pbar.set_postfix({'train_loss': train_loss / (pbar.n + 1)})

        # 验证阶段
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for img, mask in val_loader:
                img, mask = img.to(device), mask.to(device)
                pred = model(img)
                val_loss += criterion(pred, mask).item()

        val_loss /= len(val_loader)
        print(f'Validation Loss: {val_loss:.4f}')

        # 保存最佳模型
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_model.pth')
            print(f"Saved best model with val_loss: {best_val_loss:.4f}")

        scheduler.step(val_loss)

    return model


# ------------------------- 可视化 -------------------------
def visualize_results(model, loader, device, threshold=0.5):
    model.eval()
    fig, axes = plt.subplots(3, 3, figsize=(12, 12))

    with torch.no_grad():
        for i, (img, mask) in enumerate(loader):
            if i >= 3:
                break

            img, mask = img.to(device), mask.to(device)
            pred = model(img)

            # 转换为numpy
            img_np = img.squeeze().permute(1, 2, 0).cpu().numpy()
            mask_np = mask.squeeze().cpu().numpy()
            pred_np = pred.squeeze().cpu().numpy()

            # 后处理
            binary = (pred_np > threshold).astype(np.uint8)
            kernel = np.ones((3, 3), np.uint8)
            processed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

            # 可视化
            axes[i, 0].imshow(img_np)
            axes[i, 0].set_title('Input')
            axes[i, 0].axis('off')

            axes[i, 1].imshow(mask_np, cmap='gray')
            axes[i, 1].set_title('GT Mask')
            axes[i, 1].axis('off')

            overlay = img_np.copy()
            overlay[processed == 1] = [1, 0, 0]  # 红色标记
            axes[i, 2].imshow(overlay)
            axes[i, 2].set_title('Detection')
            axes[i, 2].axis('off')

    plt.tight_layout()
    plt.show()


# ------------------------- 主函数 -------------------------
def main():
    # 参数配置
    data_dir = '/home/wc/disk1/DocTamper/DocTamperV1/DocTamperV1-TrainingSet'  # 包含tampered/和mask/子目录
    batch_size = 8
    patch_size = 512
    epochs = 50
    lr = 1e-4

    # 创建数据集
    dataset = TamperDataset(
        tampered_dir=os.path.join(data_dir, 'img'),
        mask_dir=os.path.join(data_dir, 'mask'),
        patch_size=patch_size
    )

    # 分割数据集
    train_size = int(0.9 * len(dataset))
    val_size = len(dataset) - train_size
    train_set, val_set = torch.utils.data.random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=4)

    # 创建模型
    model = make_net(
        nplanes_in=3,
        kernels=[3] * 15,
        features=[64] * 14 + [1],
        bns=[False] + [True] * 13 + [False],
        acts=['relu'] * 14 + ['sigmoid'],
        dilats=[1] * 15,
        bn_momentum=0.1,
        padding=1
    )

    # 训练
    trained_model = train_model(model, train_loader, val_loader, epochs=epochs, lr=lr)

    # 可视化
    visualize_results(trained_model, val_loader, torch.device('cuda' if torch.cuda.is_available() else 'cpu'))


if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    main()